package com.zhouyu.service;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.zhouyu.domain.User;
import com.zhouyu.entity.DemoData;
import com.zhouyu.mapper.UserMapper;
import com.zhouyu.utils.TestFileUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * @author lipengze
 * @version 1.0
 * @description: TODO
 * @date 2024/11/28 13:31
 */
@Service
@Slf4j
public class PoolZipService {

    @Autowired
    private UserMapper userMapper;

    private static final int THREAD_COUNT = Runtime.getRuntime().availableProcessors(); // 根据CPU核心数设置线程池大小
    private static final int PAGE_SIZE = 500000; // 每个Sheet的数据量
    private static final int TOTAL_SHEETS = 10; // 总共要导出的Sheet数量

    private final Object writeLock = new Object(); // 写入Excel时使用的锁对象，确保多线程安全

    /**
     * 导出Excel文件并压缩为ZIP包
     * @return 压缩后的ZIP文件的输入流
     * @throws IOException 如果文件操作过程中出现异常
     */
    public FileInputStream exportToZip() throws IOException {
        long startTime = System.currentTimeMillis(); // 记录开始时间

        // 生成Excel文件的路径
        String excelFileName = TestFileUtil.getPath() + "repeatedWrite" + System.currentTimeMillis() + ".xlsx";
        // 生成ZIP文件的路径
        String zipFileName = TestFileUtil.getPath() + "exportedData" + System.currentTimeMillis() + ".zip";

        try (ExcelWriter excelWriter = EasyExcel.write(excelFileName, DemoData.class).build()) { // 创建ExcelWriter对象
            ExecutorService executorService = Executors.newFixedThreadPool(THREAD_COUNT); // 创建固定大小的线程池
            CountDownLatch countDownLatch = new CountDownLatch(TOTAL_SHEETS); // 创建CountDownLatch对象，用于等待所有线程完成

            for (int i = 1; i <= TOTAL_SHEETS; i++) { // 遍历所有Sheet
                int finalI = i; // 因为lambda表达式需要的是有效的局部变量，所以这里创建一个final变量
                executorService.submit(() -> { // 提交任务到线程池
                    try {
                        log.info("正在处理第{}页", finalI); // 日志：正在处理的页码
                        WriteSheet writeSheet = EasyExcel.writerSheet(finalI - 1, "第" + finalI  + "页").build(); // 创建WriteSheet对象，设置Sheet名称
                        List<DemoData> data = fetchDataForSheet(finalI, PAGE_SIZE); // 获取该Sheet的数据
                        synchronized (writeLock) { // 同步块，确保多线程下写入操作的安全性
                            excelWriter.write(data, writeSheet); // 将数据写入Excel
                        }
                    } catch (Exception e) {
                        log.error("第{}页导出失败", finalI, e); // 如果发生异常，记录错误日志
                    } finally {
                        countDownLatch.countDown(); // 无论成功还是失败，都减少CountDownLatch计数
                    }
                });
            }

            countDownLatch.await(); // 等待所有线程完成
            executorService.shutdown(); // 关闭线程池
        } catch (Exception e) {
            log.error("导出失败", e); // 记录导出过程中的异常
            throw new RuntimeException("导出失败", e);
        }

        // 压缩Excel文件到ZIP
        File zipFile = createZipFile(excelFileName, zipFileName);

        long endTime = System.currentTimeMillis(); // 记录结束时间
        log.info("导出成功，耗时: {} 毫秒", endTime - startTime); // 输出导出总耗时

        // 删除原始Excel文件
        new File(excelFileName).delete();

        // 返回压缩文件的输入流
        return new FileInputStream(zipFile);
    }

    /**
     * 创建ZIP文件
     * @param excelFile Excel文件的路径
     * @param zipFile ZIP文件的路径
     * @return 压缩后的文件对象
     * @throws IOException 如果文件操作过程中出现异常
     */
    private File createZipFile(String excelFile, String zipFile) throws IOException {
        File file = new File(zipFile);
        try (FileOutputStream fos = new FileOutputStream(file);
             ZipOutputStream zos = new ZipOutputStream(fos)) {
            ZipEntry entry = new ZipEntry("data.xlsx"); // 创建ZIP条目
            zos.putNextEntry(entry); // 添加条目到ZIP文件
            try (FileInputStream fis = new FileInputStream(new File(excelFile))) {
                byte[] buffer = new byte[1024]; // 缓冲区
                int length;
                while ((length = fis.read(buffer)) > 0) {
                    zos.write(buffer, 0, length); // 将Excel文件的内容写入ZIP条目
                }
            }
            zos.closeEntry(); // 关闭当前条目
        }
        return file;
    }

    /**
     * 从数据库中查询指定Sheet的数据
     * @param sheetIndex  Sheet索引（从1开始）
     * @param pageSize  每个Sheet的数据量
     * @return  分页数据列表
     */
    private List<DemoData> fetchDataForSheet(int sheetIndex, int pageSize) {
        int startPage = (sheetIndex - 1) * (pageSize / 50000) + 1; // 计算开始页码
        int endPage = startPage + (pageSize / 50000) - 1; // 计算结束页码

        List<DemoData> demoDataList = new ArrayList<>(); // 创建列表存储查询结果

        for (int page = startPage; page <= endPage; page++) { // 循环遍历每一页
            Page<User> pageObj = new Page<>(page, 50000); // 创建分页对象，每次查询50000条数据
            IPage<User> pageResult = userMapper.selectPage(pageObj, null); // 查询数据库
            for (User user : pageResult.getRecords()) { // 遍历查询结果
                DemoData demoData = new DemoData(); // 创建DemoData对象
                demoData.setId(user.getId()); // 设置ID
                demoData.setName(user.getName() != null ? user.getName() : ""); // 设置姓名，如果为空则设置为空字符串
                demoData.setAge(user.getAge() != null ? user.getAge() : 0); // 设置年龄，如果为空则设置为0
                demoData.setSex(user.getSex() != null ? user.getSex() : ""); // 设置性别，如果为空则设置为空字符串
                demoDataList.add(demoData); // 添加到列表中
            }
        }
        return demoDataList; // 返回查询结果列表
    }
}